#include "yolov8-seg_deploy.h"

yolov8_seg_deploy::yolov8_seg_deploy(const std::string& modelPath, const std::string& classPath) : 
    trtModelDeploy(modelPath), 
    classLabelPath(classPath)
{
   init_size(classLabelPath);  
}   

void yolov8_seg_deploy::infer(const cv::Mat& image)       
{
    memcpy(input, image.ptr<float>(), inputsize * sizeof(float));
    for(int i = 0; i < 100; i++)
    {
        start_time = cv::getTickCount();
        cudaMemcpyAsync(buffers[0], input, inputsize*sizeof(float), cudaMemcpyHostToDevice, stream);
        m_context->enqueueV2(buffers, stream, nullptr);
        cudaStreamSynchronize(stream);
        cudaMemcpyAsync(output1, buffers[1], outputsize1*sizeof(float), cudaMemcpyDeviceToHost, stream);
        cudaMemcpyAsync(output, buffers[2], outputsize*sizeof(float), cudaMemcpyDeviceToHost, stream);
        cudaStreamSynchronize(stream);
        end_time = cv::getTickCount();

        float t = (end_time - start_time) / static_cast<float>(cv::getTickFrequency());
        std::cout << "the fps is " << 1.0/t << std::endl;
    }
}           
  
void yolov8_seg_deploy::init_size(const std::string& classPath)
{
    
    labels = TRTHelper::getInstance().readClassNames(classPath);
     
    int32_t num =  m_engine->getNbBindings();

    for(int i = 0; i < num; i++)
    {
        std::cout << "the name is " << m_engine->getBindingName(i) << std::endl;
        std::string layname = m_engine->getBindingName(i);
        Dims shapeDims = m_engine->getBindingDimensions(i);
        if(layname == "images")
        {
            for (int j = 0; j < shapeDims.nbDims; j++)
            {
                std::cout << shapeDims.d[j];
                inputsize *= shapeDims.d[j];
                if(j == 2)
                {
                    input_h = shapeDims.d[j];
                }else if(j == 3)
                {
                    input_w = shapeDims.d[j];
                }
                if(j == shapeDims.nbDims - 1)
                    break;
                std::cout << "x";
            }
            std::cout << std::endl;
        }else if(layname == "output0")
        {
            for (int j = 0; j < shapeDims.nbDims; j++)
            {
                std::cout << shapeDims.d[j];   // 打印出张量维度
                outputsize *= shapeDims.d[j];
                if(j == 1)
                {
                    out_num = shapeDims.d[j];
                }else if(j == 2)
                {
                    out_ch = shapeDims.d[j];
                }
                if(j == shapeDims.nbDims - 1)
                    break;
                std::cout << "x";
            }
            std::cout << std::endl;
        }else if(layname == "output1")
        {
            for (int j = 0; j < shapeDims.nbDims; j++)
            {
                std::cout << shapeDims.d[j];   // 打印出张量维度
                outputsize1 *= shapeDims.d[j];
                if(j == shapeDims.nbDims - 1)
                    break;
                std::cout << "x";
            }
            std::cout << std::endl;
        }
    }
    std::cout << "input_h is " << input_h << " input_w is " << input_w << std::endl;
    std::cout << "inputSize is " << inputsize << std::endl;
    std::cout << "outputSize is " << outputsize << std::endl;
    std::cout << "outputsize1 is " << outputsize1 << std::endl;

    sx = 160.0f / 640.0f;
    sy = 160.0f / 640.0f;

    // cudaMallocHost
    cudaMallocHost(&input, inputsize*sizeof(float));
    cudaMallocHost(&output, outputsize * sizeof(float));
    cudaMallocHost(&output1, outputsize1 * sizeof(float));

    cudaMalloc(&buffers[0], inputsize * sizeof(float));
    cudaMalloc(&buffers[1], outputsize1 * sizeof(float));
    cudaMalloc(&buffers[2], outputsize * sizeof(float));
}

cv::Mat yolov8_seg_deploy::preprocess(const cv::Mat& image)
{
   int w = image.cols;
    int h = image.rows;

    int _max = std::max(h,w);

    cv::Mat image_m = cv::Mat::zeros(cv::Size(_max, _max), CV_8UC3);
    cv::Rect roi(0,0,w,h);
    image.copyTo(image_m(roi));
    x_factor = image_m.cols / static_cast<float>(input_h);
    y_factor = image_m.rows / static_cast<float>(input_w);

    cv::Mat blob = cv::dnn::blobFromImage(image_m, 1.0/255.0, cv::Size(input_w, input_h),
                                          cv::Scalar(0,0,0), true, true);

    return blob;
}

void yolov8_seg_deploy::postprocess(const cv::Mat& image)
{
    float *pdata = output;
    float *mdata = output1;

    // 后处理 1x25200x85 85-box conf 80- min/max
    std::vector<cv::Rect> boxes;
    std::vector<int> classIds;
    std::vector<float> confidences;

    std::vector<cv::Mat> masks;
    cv::Mat mask1(32, 25600, CV_32F, (float*)mdata);

    cv::Mat det_output(out_num, out_ch, CV_32F, (float*)pdata);
    det_output = det_output.t();


    for(int i = 0; i < det_output.rows; i++)
    { 
        cv::Mat classes_scores = det_output.row(i).colRange(4,  out_num - 32);
        cv::Point classIdPoint;
        double score;
        cv::minMaxLoc(classes_scores, 0, &score, 0, &classIdPoint);

        // 置信度0-1之间
        if( score > 0.25)
        {
            cv::Mat mask2 = det_output.row(i).colRange(out_num - 32,  out_num);
            float cx = det_output.at<float>(i,0);
            float cy = det_output.at<float>(i,1);
            float ow = det_output.at<float>(i,2);
            float oh = det_output.at<float>(i,3);

            int x = static_cast<int>((cx - 0.5*ow) * x_factor);
            int y = static_cast<int>((cy - 0.5*oh) * y_factor);
            int width = static_cast<int>(ow*x_factor);
            int height = static_cast<int>(oh*y_factor);

            cv::Rect box;
            box.x = x;
            box.y = y;
            box.width = width;
            box.height = height;

            boxes.push_back(box);
            classIds.push_back(classIdPoint.x);
            confidences.push_back(score);
            masks.push_back(mask2);
        }
    }

    // NMS
    std::vector<int> indexes;
    cv::dnn::NMSBoxes(boxes, confidences, (float)(0.25), (float)(0.45), indexes);
    cv::Mat rgb_mask = cv::Mat::zeros(image.size(), image.type());
    for(size_t i = 0; i < indexes.size(); i++)
    {
        int idx = indexes[i];
        int cid = classIds[idx];

        cv::Rect box = boxes[idx];
        int x1 = std::max(0, box.x);
        int y1 = std::max(0, box.y);
        int x2 = std::max(0, box.br().x);
        int y2 = std::max(0, box.br().y);
        cv::Mat m2 = masks[idx];
        cv::Mat m = m2 * mask1;
        for (int col = 0; col < m.cols; col++) {
            m.at<float>(0, col) =  TRTHelper::getInstance().sigmoid_function(m.at<float>(0, col));
        }
        cv::Mat m1 = m.reshape(1, 160);
        int mx1 = std::max(0, int((x1 * sx) / x_factor));
        int mx2 = std::max(0, int((x2 * sx) / x_factor));
        int my1 = std::max(0, int((y1 * sy) / y_factor));
        int my2 = std::max(0, int((y2 * sy) / y_factor));
        // fix out of range box boundary on 2022-12-14
        if (mx2 >= m1.cols) {
            mx2 = m1.cols - 1;
        }
        if (my2 >= m1.rows) {
            my2 = m1.rows - 1;
        }
        // end fix it!!

        cv::Mat mask_roi = m1(cv::Range(my1, my2), cv::Range(mx1, mx2));
        cv::Mat rm, det_mask;
        cv::resize(mask_roi, rm, cv::Size(x2 - x1, y2 - y1));
        for (int r = 0; r < rm.rows; r++) {
            for (int c = 0; c < rm.cols; c++) {
                float pv = rm.at<float>(r, c);
                if (pv > 0.5) {
                    rm.at<float>(r, c) = 1.0;
                }
                else {
                    rm.at<float>(r, c) = 0.0;
                }
            }
        }
        rm = rm * rng.uniform(0, 255);
        rm.convertTo(det_mask, CV_8UC1);
        if ((y1 + det_mask.rows) >= image.rows) {
            y2 = image.rows - 1;
        }
        if ((x1 + det_mask.cols) >= image.cols) {
            x2 = image.cols - 1;
        }
        // std::cout << "x1: " << x1 << " x2:" << x2 << " y1: " << y1 << " y2: " << y2 << std::endl;
        cv::Mat mask = cv::Mat::zeros(cv::Size(image.cols, image.rows), CV_8UC1);
        det_mask(cv::Range(0, y2 - y1), cv::Range(0, x2 - x1)).copyTo(mask(cv::Range(y1, y2), cv::Range(x1, x2)));
        add(rgb_mask, cv::Scalar(rng.uniform(0, 255), rng.uniform(0, 255), rng.uniform(0, 255)), rgb_mask, mask);

        cv::rectangle(image, boxes[idx], cv::Scalar(0,0,255), 2, 8,0);
        cv::putText(image, cv::format("%s_%.2f", labels[cid].c_str(), confidences[idx]) , boxes[idx].tl(),
                    cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(0,255,0), 2, 8);
    }

    // compute the fps
    float t = (end_time - start_time) / static_cast<float>(cv::getTickFrequency());
    cv::putText(image, cv::format("FPS: %.2f", 1.0/t), cv::Point(20,40), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(255, 0, 0), 2, 8);

    cv::Mat result;
   cv::addWeighted(image, 0.5, rgb_mask, 0.5, 0, result);
   result.copyTo(image);
    cv::imwrite("./test-yolov8-seg.jpg", image);
}
